package com.gjxjl.jdbc.gateway.handler;

import com.gjxjl.jdbc.gateway.config.ProxyConfig;
import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.sql.*;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class JdbcProxyHandler extends ChannelInboundHandlerAdapter {
    private static final Logger logger = LogManager.getLogger(JdbcProxyHandler.class);

    private final ProxyConfig config;
    private final Map<Channel, Connection> clientConnections;
    private Connection backendConnection;

    public JdbcProxyHandler(ProxyConfig config) {
        this.config = config;
        this.clientConnections = new ConcurrentHashMap<>();
        initializeBackendConnection();
    }

    private void initializeBackendConnection() {
        try {
            String url = String.format("jdbc:mysql://%s:%d/%s",
                    config.getBackendHost(),
                    config.getBackendPort(),
                    config.getDefaultDatabase());

            this.backendConnection = DriverManager.getConnection(
                    url, config.getBackendUsername(), config.getBackendPassword());

            logger.info("Backend database connection established");
        } catch (SQLException e) {
            logger.error("Failed to connect to backend database", e);
            throw new RuntimeException("Backend connection failed", e);
        }
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        if (msg instanceof ByteBuf) {
            ByteBuf buf = (ByteBuf) msg;
            try {
                // 解析JDBC协议数据包
                byte[] data = new byte[buf.readableBytes()];
                buf.readBytes(data);

                // 处理JDBC请求
                processJdbcRequest(ctx, data);
            } finally {
                buf.release();
            }
        }
    }

    private void processJdbcRequest(ChannelHandlerContext ctx, byte[] requestData) {
        try {
            // 这里简化处理，实际需要解析MySQL协议
            String sql = parseSqlFromRequest(requestData);

            if (config.isEnableLogging()) {
                logger.info("SQL Executed: {}", sql);
            }

            // 安全检查
            if (isSqlBlocked(sql)) {
                sendErrorResponse(ctx, "SQL blocked by security policy");
                return;
            }

            // 执行SQL并获取结果
            executeSqlAndSendResponse(ctx, sql);

        } catch (Exception e) {
            logger.error("Error processing JDBC request", e);
            sendErrorResponse(ctx, e.getMessage());
        }
    }

    private String parseSqlFromRequest(byte[] data) {
        // 简化实现，实际需要解析MySQL协议
        return new String(data).trim();
    }

    private boolean isSqlBlocked(String sql) {
        return config.getBlockedSqlPatterns().stream()
                .anyMatch(pattern -> sql.toLowerCase().contains(pattern.toLowerCase()));
    }

    private void executeSqlAndSendResponse(ChannelHandlerContext ctx, String sql) {
        try (Statement stmt = backendConnection.createStatement()) {
            long startTime = System.currentTimeMillis();

            if (sql.trim().toLowerCase().startsWith("select")) {
                // 查询操作
                ResultSet rs = stmt.executeQuery(sql);
                sendQueryResponse(ctx, rs);
            } else {
                // 更新操作
                int affectedRows = stmt.executeUpdate(sql);
                sendUpdateResponse(ctx, affectedRows);
            }

            long endTime = System.currentTimeMillis();

            if (config.isEnableMetrics()) {
                logger.info("SQL Execution Time: {}ms", (endTime - startTime));
            }

        } catch (SQLException e) {
            logger.error("SQL execution error", e);
            sendErrorResponse(ctx, e.getMessage());
        }
    }

    private void sendQueryResponse(ChannelHandlerContext ctx, ResultSet rs) throws SQLException {
        // 简化实现，实际需要构建MySQL协议响应包
        StringBuilder response = new StringBuilder();
        ResultSetMetaData metaData = rs.getMetaData();
        int columnCount = metaData.getColumnCount();

        // 构建结果集
        while (rs.next()) {
            for (int i = 1; i <= columnCount; i++) {
                response.append(rs.getString(i)).append("\t");
            }
            response.append("\n");
        }

        ByteBuf buf = ctx.alloc().buffer(response.length());
        buf.writeBytes(response.toString().getBytes());
        ctx.writeAndFlush(buf);
    }

    private void sendUpdateResponse(ChannelHandlerContext ctx, int affectedRows) {
        String response = "Affected rows: " + affectedRows;
        ByteBuf buf = ctx.alloc().buffer(response.length());
        buf.writeBytes(response.getBytes());
        ctx.writeAndFlush(buf);
    }

    private void sendErrorResponse(ChannelHandlerContext ctx, String error) {
        String response = "ERROR: " + error;
        ByteBuf buf = ctx.alloc().buffer(response.length());
        buf.writeBytes(response.getBytes());
        ctx.writeAndFlush(buf);
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        logger.error("Channel exception", cause);
        ctx.close();
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        clientConnections.remove(ctx.channel());
        super.channelInactive(ctx);
    }
}
